import copy

from flask import g
import bson
from bson import ObjectId
from wtforms import fields, validators, ValidationError
from wtforms.validators import DataRequired

from app.db.models.assets import Asset
from app.libs.ip_local_info import get_ip_location
from app.libs.score import Score
from app.validators import BaseForm, check_object_id, UidForm, PageForm
from app.db.models.tasks import Task
from app.db.models.jobs import Job
from app.db.models.wass import ScanEvents
from app.errors import JobNotFound, TaskNotFound, AssetNotFound
from app.libs.enums import (
    FOREIGN_LINKS_PLUGIN_TUPLE,
    WS_TASK_GROUP_TUPLE,
    TASK_TYPE_CONFIG
)
from app.libs.utils import get_used_time
from app.errors import ParameterError
from app.db.models.asset_tasks import AssetTasks, get_monitor_urls
from app.config.settings import SPIDER_URL_TIME


class TaskSummaryForm(BaseForm):
    uid = fields.StringField()

    def validate_uid(self, value):
        if uid := value.data:
            try:
                value.data = ObjectId(uid)
            except Exception:
                raise ValidationError('uid format error')
        if g.entity_type == 'user':
            value.data = g.user.pk

    def _validate_object_id(self, data):
        try:
            data = ObjectId(data)
        except Exception as e:
            raise ValidationError(e)

        return data


class TaskDetailForm(TaskSummaryForm):
    searchKey = fields.StringField()
    exSearchKey = fields.StringField()
    page = fields.IntegerField(default=1)
    perPage = fields.IntegerField(default=10, validators=[validators.NumberRange(max=100)])


class HttpTaskDetailForm(TaskDetailForm):
    securityStatus = fields.StringField(default='')


class VulTaskDetailForm(TaskDetailForm):
    securityStatus = fields.StringField(default='')
    triggerType = fields.StringField()


class SslTaskDetailForm(VulTaskDetailForm):
    pass


class SecurityEventTaskDetailForm(VulTaskDetailForm):
    pass


class ContentTaskDetailForm(VulTaskDetailForm):
    pass


class AssetTaskDetailForm(BaseForm):
    job_id = fields.StringField()
    uid = fields.StringField()
    task_id = fields.StringField()
    status = fields.BooleanField(default=False)
    warning_ids = fields.FieldList(fields.StringField())

    def validate_job_id(self, value):
        if job_id := value.data:
            try:
                value.data = ObjectId(job_id)
            except Exception:
                raise ValidationError('job_id format error')

    def validate_uid(self, value):
        if uid := value.data:
            try:
                value.data = ObjectId(uid)
            except Exception:
                raise ValidationError('uid format error')

    def validate_task_id(self, value):
        if task_id := value.data:
            try:
                value.data = ObjectId(task_id)
            except Exception:
                raise ValidationError('task_id format error')

    @staticmethod
    def get_detail_overview(uid, job_id, task_id=None):
        job = Job.objects.filter(uid=uid, id=job_id).first()
        if not job:
            raise JobNotFound()

        if task_id:
            task = Task.objects.filter(id=task_id, taskType="asset", status="completed").first()
            if not task:
                return {}
        else:
            asset = Asset.objects.filter(id=job.assetId).first()
            if not asset:
                raise AssetNotFound()

            if asset.resultId:
                task = Task.objects.filter(id=asset.resultId).first()
                if not task:
                    return {}
            else:
                return {}

        addition = task.result.addition
        icp = addition.get("icp", {"desc": "--", "icp": "", "status": 2})
        if icp.get("desc") not in ("未备案", "已备案"):
            icp["desc"] = "--"

        source_ip = addition.get("port", {}).get("ip", "")
        region = get_ip_location(source_ip, country=True, province=True, city=True, res_type="list")
        ipdb = {"country_name": "", "region_name": "", "city_name": ""}
        if region != "未知":
            ipdb = {"country_name": region[0], "region_name": region[1], "city_name": region[2]}

        result = {
            "target_ip": addition.get("target_ip", ""),
            "waf_list": addition.get("waf_list", []),
            "alexa": addition.get("alexa", ''),
            "basic_info": addition.get("basic_info", {}),
            "icp": icp,
            "ipdb": ipdb,
            "http_get": addition.get("http_get", {}),
            "http_method": addition.get("http_method", ""),
            "port": addition.get("port", {'ip': '', 'host': '', 'data': []}),
        }

        subdomain = addition.get('subdomain', [])
        all_target_set = set()
        subdomain_list = []
        for item in subdomain:
            for ip in item.get("ip").split(","):
                _item = copy.deepcopy(item)
                _item['ip'] = ip
                subdomain_list.append(_item)

            item["accessed"] = False
            all_target_set.add(item.get("subdomain"))

        exist_jobs = set()
        if all_target_set:
            jobs = Job.objects.find({"uid": uid, "target.host": {"$in": list(all_target_set)}})
            for job in jobs:
                source_ip = job.sourceIp if job.sourceIp else ""
                exist_jobs.add(f"{job.target.host}{source_ip}")

        if exist_jobs:
            for item in subdomain_list:
                key = f'{item.get("subdomain")}{item.get("ip")}'
                if key in exist_jobs:
                    item["accessed"] = True

        result["subdomain"] = subdomain_list

        return result

    @staticmethod
    def get_detail_change(uid, job_id, task_id=None):
        job = Job.objects.filter(uid=uid, id=job_id).first()
        if not job:
            raise JobNotFound()

        result = {
            "is_last_task": False,
            "subdomain": {"add": [], "remove": []},
            "port": {"add": [], "remove": []},
        }

        asset = Asset.objects.filter(id=job.assetId).first()
        if not asset:
            raise AssetNotFound()

        if task_id:
            task = Task.objects.filter(id=task_id, taskType="asset", status="completed").first()
            if not task:
                return result
            if asset.resultId and str(task_id) == str(asset.resultId):
                result['is_last_task'] = True
        else:
            if asset.resultId:
                task = Task.objects.filter(id=asset.resultId).first()
                if not task:
                    return result
            else:
                return result

            result['is_last_task'] = True

        for item in task.result.warnings:
            category = item.category
            if category in ('subdomain', 'port'):
                if category == 'port':
                    try:
                        item.affects = int(float(item.affects))
                    except Exception:
                        pass
                result[category][item.detail.get('op')].append(item)

        return result

    @staticmethod
    def get_detail_summary(uid, job_id):
        job = Job.objects.filter(uid=uid, id=job_id).first()
        if not job:
            raise JobNotFound()

        warning_count_map = {TASK_TYPE_CONFIG[t]["line"]: 0 for t in WS_TASK_GROUP_TUPLE}
        data = {
            "uid": job.uid,
            "note": job.note,
            "target_url": job.targetUrl,
            "job_id": job.id,
            "source_ip": job.sourceIp,
            "level": '未知',
            "score": 0,
            "is_verified": job.verification.isVerified,
            "warning_count_map": warning_count_map,
            "update_time": "",
            "is_checked": False,
            "is_display_url": False
        }
        schedule_task_info = {}
        asset_task = AssetTasks.objects.find({
            "jobId": job.id,
            "taskGroupType": "web_security",
            "triggerType": "schedule"
        }).first()
        if asset_task:
            schedule_task_info.update({
                "enable_monitor": asset_task.enableMonitor,
                "asset_task_id": asset_task.id
            })
            for t in WS_TASK_GROUP_TUPLE:
                _enable = getattr(asset_task, f"{t}Settings").enable
                _task_id = getattr(asset_task, f"{t}Settings").taskId
                schedule_task_info.update({
                    TASK_TYPE_CONFIG[t]["line"]:  _enable and bool(_task_id)
                })
        data.update({"schedule_web_task_info": schedule_task_info})

        query = {'_id': job.id}
        lookups = [
            {'$lookup': {'from': "tasks", 'localField': f"{taskType}ResultId", 'foreignField': "_id",
                         'as': f"{taskType}"}}
            for taskType in WS_TASK_GROUP_TUPLE
        ]
        _projects = ['{}.taskSettings.enable', '{}.securityStatus', '{}.endTime', '{}.result', '{}.startTime']
        project = {_project.format(task_type): 1 for _project in _projects for task_type in WS_TASK_GROUP_TUPLE}
        aggregate = [
            {'$match': query},
            *lookups,
            {'$project': project}
        ]
        job_status = {job['_id']: job for job in Job.objects.aggregate(aggregate)}

        result_status = job_status.get(job.id)
        if not result_status:
            return data

        level = '未知'
        score = 0
        update_time = ''
        all_warnings = []
        is_checked = False
        target_status_list = []
        task_start_times = []
        for task_type in WS_TASK_GROUP_TUPLE:
            task_type_data = result_status.get(task_type, [])
            if task_type_data:
                is_checked = True
            for item in task_type_data:
                end_time = item.get('endTime')
                task_start_times.append(item.get("startTime"))
                if not update_time or end_time > update_time:
                    update_time = end_time
                warning_count_map[TASK_TYPE_CONFIG[task_type]["line"]] = item.get('result', {}).get('warningCount', 0)

                warnings = item.get('result', {}).get('warnings', [])
                all_warnings.extend(warnings)
                target_status = item.get("result", {}).get("targetStatus", {"status": "good"})
                target_status_list.append({"end_time": item.get("endTime"), "target_status": target_status})
        if all_warnings:
            score_info = Score.score_info(event_map={}, warnings=all_warnings)
            score = score_info.get('score', 0)
            level = score_info.get('level', '安全')
        target_status_set = {i.get("target_status", {}).get("status") for i in target_status_list}
        if (not all_warnings and update_time) and (not target_status_set or "good" in target_status_set):
            level = '安全'
        elif (not all_warnings and update_time) and (target_status_set and "good" not in target_status_set):
            level = "未知"
        if target_status_list := [ts for ts in target_status_list if ts.get("target_status", {}).get("httpGetSequence")]:
            data.update({
                "target_status": list(sorted(target_status_list, key=lambda x: x["end_time"]))[-1]["target_status"]
            })
        else:
            data.update({"target_status": {"status": "good"}})
        data['score'] = score
        data['level'] = level
        data['update_time'] = update_time
        data['is_checked'] = is_checked
        data['url_count'] = len(get_monitor_urls(job_obj=job))
        data['scan_count'] = get_monitor_urls(job_obj=job, is_scan_count=True)
        data['is_display_url'] = SPIDER_URL_TIME < min(task_start_times) if task_start_times else False
        return data

    @staticmethod
    def change_resolved_status(task_id, warning_id_list, status):
        task = Task.objects.filter(id=task_id, taskType="asset", status="completed").first()
        if not task:
            raise TaskNotFound()

        result = []
        warnings = task.result.warnings
        warning_maps = dict()
        for warning in warnings:
            try:
                warning.affects = str(int(float(warning.affects)))
            except Exception:
                warning.affects = str(warning.affects)
            warning_maps[warning.id] = warning

        for warning_id in warning_id_list:
            if warning_id in warning_maps:
                if warning_maps[warning_id].isResolved == status:
                    continue
                result.append(warning_id)
                warning_maps[warning_id].isResolved = status

        if result:
            task.update(result__warnings=warnings)

        return result, task.target


class BaseEvidenceForm(TaskSummaryForm):
    warningId = fields.StringField()
    id = fields.StringField()

    def validate_warningId(self, value):
        value.data = self._validate_object_id(value.data)

    def validate_id(self, value):
        value.data = self._validate_object_id(value.data)


class OverViewForm(BaseForm):
    task_id = fields.StringField(validators=(validators.DataRequired(),))

    def validate_task_id(self, value):
        if task_id := value.data:
            try:
                value.data = ObjectId(task_id)
            except Exception:
                raise ValidationError('task_id format error')
            else:
                if g.entity_type != "app":
                    self.task = Task.objects.filter(uid=g.user.id, id=value.data).first()
                else:
                    self.task = Task.objects.filter(id=value.data).first()
                if not self.task:
                    raise ParameterError(msg="暂无数据")
                self.job = Job.objects.filter(id=self.task.jobId).first()
                if not self.job:
                    raise ParameterError(msg="暂无数据")

    def get_data(self):
        addition = self.task.result.addition or {}
        statistics = addition.get("statistics", {})
        data = {
            "id": self.task.id,
            "triggerType": self.task.triggerType,
            "taskType": self.task.taskType,
            "note": self.job.note,
            "targetUrl": self.job.targetUrl,
            "alertSettings": self.task.alertSettings,
            "taskSettings": self.task.taskSettings,
            "jobId": self.job.id,
            "endTime": self.task.endTime,
            "usedTime": self.get_used_time(),
            "wafInfo": self.get_waf_info(addition),
            "targetIp": addition.get("target_ip", "") or "暂无",
            "urlCount": statistics.get("urls", 0) if statistics else "暂无"
        }
        return data

    def get_used_time(self):
        diff_time = int((self.task.endTime - self.task.startTime).total_seconds())
        return get_used_time(diff_time)

    @staticmethod
    def get_waf_info(addition):
        waf_list = addition.get("site_info", {}).get("waf_list", [])
        if not waf_list:
            return "暂无"
        return "|".join(waf_list)


class DeleteForm(BaseForm):
    ids = fields.FieldList(fields.StringField(validators=[check_object_id]))


class HostVulTaskForm(BaseForm):
    job_id = fields.StringField(validators=(DataRequired(),))
    id = fields.StringField()
    page = fields.IntegerField(default=1)
    limit = fields.IntegerField(default=10)
    vuln_name = fields.StringField()
    vuln_type = fields.StringField()
    level = fields.IntegerField()

    def validate_job_id(self, value):
        if job_id := value.data:
            try:
                value.data = ObjectId(job_id)
            except Exception:
                raise ValidationError('job_id format error')
            else:
                job = Job.objects.filter(uid=g.user.id, id=job_id, assetType="host").first()
                if not job:
                    raise ParameterError(msg="暂无数据")
                self.job = job

    def validate_id(self, value):
        if task_id := value.data:
            try:
                _ = ObjectId(task_id)
            except Exception:
                raise ValidationError('id format error')
            else:
                self.task = Task.objects.filter(
                    uid=g.user.id, id=value.data, taskType="hostVul", status="completed").first()
                if not self.task:
                    raise ParameterError(msg="暂无数据")

    def validate_for_api(self):
        obj = super(HostVulTaskForm, self).validate_for_api()
        if not obj.id.data:
            obj.task = Task.objects.filter(
                uid=g.user.id, taskType="hostVul", status="completed", jobId=obj.job.id).order_by("-id").first()
            if not obj.task:
                raise TaskNotFound()
        if obj.task.jobId != obj.job.id:
            raise TaskNotFound()
        query = {"task_session_id": obj.task.taskSessionId, "event_name": "network_vuln"}
        obj.vul = ScanEvents.objects.find(query)
        obj.statistics = ScanEvents.objects.filter(
            task_session_id=obj.task.taskSessionId, event_name="statistics").first()
        return obj

    def query_vul(self):
        vuln_name = self.vuln_name.data
        vuln_type = self.vuln_type.data
        level = self.level.data
        query = {"task_session_id": self.task.taskSessionId, "event_name": "network_vuln"}
        if vuln_name:
            query.update({"detail.name.zh_cn": {"$regex": vuln_name}})
        if vuln_type:
            query.update({"detail.vuln_type.zh_cn": vuln_name})
        if level and level == 2:
            query.update({"detail.severity": {"$in": [0, 1, 2]}})
        elif level:
            query.update({"detail.severity": level})
        self.vul = ScanEvents.objects.find(query)
        return self

    def get_warnings(self, vul_events):
        warnings = []
        severityMap = Score.severityMap
        for event in vul_events:
            item = event.detail
            affects = item.get("host")
            warning = {
                "level": severityMap.get(item["severity"], {}).get("level"),
                "levelName": severityMap.get(item["severity"], {}).get("zh_CN"),
                "title": item.get("name", {}).get("zh_cn"),
                "category": item.get("vuln_type", [{}])[0].get("zh_cn"),
                "affects": affects,
                "detailText": item.get("extra_detail"),
                "recommendation": item.get("recommendation", {}).get("zh_cn"),
                "impact": item.get("impact", {}).get("zh_cn"),
                "description": item.get("description", {}).get("zh_cn"),
                "reference": item.get("reference"),
                "cnvd": item.get("cnvd"),
                "foundAt": item.get("found_at"),
                "traffics": item.get("traffics"),
                "payload": item.get("payload"),
                "isResolved": False
            }
            warnings.append(warning)
        return warnings

    @staticmethod
    def get_job(job_id):
        job = Job.objects.filter(uid=g.user.id, id=job_id, assetType="host").first()
        if not job:
            raise JobNotFound()
        return job


class WebTaskForm(BaseForm):
    task_group_type_query = {"taskGroupType": "web_security"}
    asset_task_id = fields.StringField(validators=(DataRequired(),))

    def validate_asset_task_id(self, value):
        uid = g.user.id
        asset_task_id = value.data
        if not asset_task_id:
            raise ParameterError(msg="参数错误")
        try:
            asset_task_id = bson.ObjectId(asset_task_id)
        except bson.errors.BSONError:
            raise ParameterError(msg="参数错误")
        query = {
            "id": asset_task_id,
        }
        query.update(self.task_group_type_query)
        if g.entity_type != "app":
            query.update({"uid": uid})
        asset_task = AssetTasks.objects.filter(**query).first()
        if not asset_task:
            raise ParameterError(msg="任务不存在")
        self.job = Job.objects.filter(id=asset_task.jobId).first()
        if not self.job:
            raise JobNotFound()
        self.asset_task = asset_task


class WebTaskTypeForm(OverViewForm):
    task_id = fields.StringField(validators=(DataRequired(),))

    def get_data(self):
        addition = self.task.result.addition or {}
        target_status = self.task.result.targetStatus
        parse_key = f"_parse_{TASK_TYPE_CONFIG[self.task.taskType]['line']}_warnings"
        if hasattr(AssetTasks, parse_key):
            score_info = getattr(AssetTasks, parse_key)(self.task.result.warnings)
        else:
            score_info = {}
        if target_status and target_status.status != "good" and not self.task.result.warnings:
            score_info.update({"level": "未知", "score": 0})
        elif ((not target_status) or (target_status and target_status.status == "good")) and score_info:
            score = score_info.get("score", 0)
            if score == 0:
                level = "安全"
            else:
                level = [_d["front_CN"] for _, _d in Score.severityMap.items() if _d["score_range"][0] <= score <= _d["score_range"][1]][0]
            score_info.update({"level": level})
        _warnings = [w.to_mongo().to_dict() for w in self.task.result.warnings]
        task_settings = self.task.taskSettings.to_dict()
        security_event_plugins = task_settings.get("collect", {}).get("plugins", [])
        if security_event_plugins and "foreign_links" in security_event_plugins:
            security_event_plugins.remove("foreign_links")
            security_event_plugins.extend(list(FOREIGN_LINKS_PLUGIN_TUPLE))
        data = {
            "overview": {
                "trigger_type": self.task.triggerType,
                "task_type": self.task.taskType,
                "job_name": self.job.note,
                "target_url": self.job.targetUrl,
                "waf_info": self.get_waf_info(addition),
                "target_ip": addition.get("target_ip", "") or "暂无",
                "url_count": len(get_monitor_urls(task_session_id=self.task.taskSessionId)),
                "scan_count": get_monitor_urls(task_session_id=self.task.taskSessionId, is_scan_count=True),
                "name": self.task.name,
                "is_display_url": SPIDER_URL_TIME < self.task.startTime if self.task.startTime else False
            },
            "score_info": score_info,
            "warnings": sorted(_warnings, key=lambda x: x["level"], reverse=True),
            "addition": self.task.result.addition,
            "time_info": {
                "end_time": self.task.endTime,
                "start_time": self.task.startTime,
                "used_time": self.get_used_time(),
            },
            "task_id": self.task.id,
            "alert_settings": self.task.alertSettings,
            "task_settings": task_settings,
            "job_id": self.job.id,
            "task_session_id": self.task.taskSessionId,
            "cydEnable": self.job.cydEnabled,
            "isLastTask": self.task.isLatestTask,
            "target_status": target_status or {"status": "good"},
            "proxy_ip": self.task.proxyIp
        }
        return data


class UpdateTaskWarningForm(OverViewForm):
    # TODO 废弃

    warning_ids = fields.FieldList(fields.StringField(validators=(DataRequired(),)))
    status = fields.StringField(validators=(DataRequired(),))

    def validate_status(self, value):
        if value.data not in ["ignore", "resolve", "unresolve"]:
            raise ParameterError(msg="漏洞状态设置错误")


class AvailableTaskForm(WebTaskForm):
    task_group_type_query = {"taskGroupType": {"$in": ["http", "ping"]}}


class ScheduleRunningForm(BaseForm):

    asset_task_ids = fields.FieldList(fields.StringField())

    def validate_asset_task_ids(self, value):
        asset_task_ids = value.data
        atids = []
        for atid in asset_task_ids:
            if not atid:
                continue
            try:
                atids.append(bson.ObjectId(atid))
            except bson.errors.BSONError:
                raise ParameterError(msg="任务ID格式错误")
        if not atids:
            raise ParameterError(msg="没有选择任何周期任务")
        if g.role.name == '管理员':
            self.asset_tasks = AssetTasks.objects.filter(triggerType='schedule', id__in=atids)
        else:
            self.asset_tasks = AssetTasks.objects.filter(triggerType='schedule', id__in=atids, uid=g.user.id)
        if not self.asset_tasks:
            raise ParameterError(msg="没有选择任何周期任务")


class HostTaskForm(WebTaskForm):
    task_group_type_query = {"taskGroupType": "host"}


class MonitorUrlsForm(UidForm):
    query_type = fields.StringField()
    query_id = fields.StringField()

    def validate_id(self, query_id):
        try:
            bson.ObjectId(query_id.data)
        except bson.errors.BSONError:
            raise ParameterError(msg="查询监测url, id错误")

    def validate_query_type(self, query_type):
        if query_type.data not in ["job", "task", "asset_task"]:
            raise ParameterError(msg="查询监测url, 类型错误")


class PostDelWarningIgnore(BaseForm):
    task_id = fields.StringField()
    warning_ids = fields.FieldList(fields.StringField())
    reason = fields.IntegerField()
    types = fields.FieldList(fields.StringField())

    def validate_reason(self, reason):
        if not reason.data:
            return
        if reason.data not in [1, 2, 3]:
            raise ParameterError(msg="忽略原因错误")

    def validate_task_id(self, task_id):
        if not task_id.data:
            return
        try:
            self.task_id.data = bson.ObjectId(task_id.data)
        except bson.errors.BSONError:
            raise ParameterError(msg="任务ID格式错误")

    def validate_warning_ids(self, warning_ids):
        if not warning_ids.data:
            raise ParameterError(msg="缺少威胁ID")
        try:
            self.warning_obj_ids = [bson.ObjectId(i) for i in warning_ids.data]
        except bson.errors.BSONError:
            raise ParameterError(msg="威胁ID格式错误")


class QueryWarningIgnore(PageForm):
    username = fields.StringField()
    warning_id = fields.StringField()
    level = fields.IntegerField()
    task_type = fields.StringField()
    target_url = fields.StringField()
    target_name = fields.StringField()
    risk_url = fields.StringField()
    risk_name = fields.StringField()
    reason = fields.IntegerField()
    sort = fields.StringField(default="-ignoredTime")

    def validate_warning_id(self, warning_id):
        if not warning_id.data:
            return
        try:
            self.warning_id.data = bson.ObjectId(warning_id.data)
        except bson.errors.BSONError:
            raise ParameterError(msg="威胁ID格式错误")

    def validate_sort(self, sort):
        if sort.data not in ["-ignoredTime", "ignoredTime"]:
            self.updated_at.data = "-ignoredTime"
